import numpy as np
from numba import njit
import pybroker
from pybroker import Strategy
from pybroker.ext.data import AKShare
pybroker.enable_data_source_cache('AKShare')

# 自定义指标函数，计算收盘价减去移动平均值（CMMA），可用于 均值回归 策略：
# 如果CMMA值为正，那么当前价格高于其移动平均值；如果CMMA值为负，那么当前价格低于其移动平均值。
def cmma(bar_data, lookback):

    @njit  # Enable Numba JIT.
    def vec_cmma(values):
        # Initialize the result array.
        n = len(values)
        out = np.array([np.nan for _ in range(n)])

        # For all bars starting at lookback:
        for i in range(lookback, n):
            # Calculate the moving average for the lookback.
            ma = 0
            for j in range(i - lookback, i):
                ma += values[j]
            ma /= lookback
            # Subtract the moving average from value.
            out[i] = values[i] - ma
        return out

    # Calculate with close prices.
    return vec_cmma(bar_data.close)
# 注册函数到内部
# 我们将指标函数命名为 cmma_20，并将 回溯 参数指定为 20 条。
cmma_20 = pybroker.indicator('cmma_20', cmma, lookback=20)
'''指标函数中在 bar_data 之后的任何参数都将作为用户定义的参数传递给 pybroker.indicator。
一旦指标函数在 PyBroker 中注册，它将返回一个新的 Indicator 实例，该实例引用我们定义的指标函数。'''

aKShare = AKShare()
df = aKShare.query('000001', '4/1/2020', '4/1/2022')

def buy_cmma_cross(ctx):
    if ctx.long_pos():
        return
    # Place a buy order if the most recent value of the 20 day CMMA is < 0:
    # 这里调用了 ctx 对象的 indicator 方法，并传入了一个20日CMMA（自定义移动平均差异）指标的名字。
    # [-1] 表示获取这个指标数组的最后一个值（即最新的值）。
    # 如果这个值小于0，那么说明某种交易条件被触发了（可能是CMMA线向下穿越了零线）。
    if ctx.indicator('cmma_20')[-1] < 0:
        ctx.buy_shares = ctx.calc_target_shares(1)
        ctx.hold_bars = 3 # 表示策略计划持有这些股票3个交易日

strategy = Strategy(aKShare, '4/1/2020', '4/1/2022')
strategy.add_execution(buy_cmma_cross, '000001', indicators=cmma_20)
pybroker.enable_indicator_cache('my_indicators')

result = strategy.backtest(warmup=20) #warmup 参数指定在运行回测执行之前需要经过 20 个 Bar：
# 在Python的pandas库中，DataFrame 对象有一个方法 round()，
# 该方法用于将DataFrame中的所有数值型列（通常是整数或浮点数）四舍五入到指定的小数位数。
result.metrics_df.round(4)
print(result.metrics_df.round(4))